package org.luckyjourney;

import com.baomidou.mybatisplus.core.toolkit.ObjectUtils;

import java.util.*;

public class ProbabilityArrayInitializer {
    /**
     * @param modelMap 用户模型  标签及标签权重
     * @param cnt      随机的次数
     * @return List<String>  根据权重抽取到的标签列表
     */
    public List<String> initProbabilityArray(Map<String, Double> modelMap, int cnt) {
        if (ObjectUtils.isEmpty(modelMap)) {
            return new ArrayList<>();
        }
        // 获取到所有标签
        List<String> labels = new ArrayList<>(modelMap.size());
        List<Double> weight = new ArrayList<>(modelMap.size());
        for (Map.Entry<String, Double> entry : modelMap.entrySet()) {
            // 将标签放入集合中
            labels.add(entry.getKey());
            // 将权重放入集合中
            weight.add(entry.getValue());
        }

        // 初始化权重分布
        double[] probability = new double[modelMap.size()];
        probability[0] = weight.get(0);
        for (int i = 1; i < weight.size(); i++) {
            probability[i] = probability[i - 1] + weight.get(i);
        }
        // 获取总权重
        double sum = probability[probability.length - 1];

        // 随机count次
        Random random = new Random();
        // 存储随机到哪些标签
        List<String> res = new ArrayList<>();
        for (int i = 0; i < cnt; i++) {
            double r = random.nextDouble() * sum;
            // 计算落在哪个区间里
            for (int j = 0; j < probability.length; j++) {
                if (r < probability[j]) {
                    res.add(labels.get(j));
                    break;
                }
            }
        }
        return res;
    }

    public static void main(String[] args) {
        Map<String, Double> modelMap = new HashMap<>();
        modelMap.put("A", 0.1);
        modelMap.put("B", 0.2);
        modelMap.put("C", 0.3);
        modelMap.put("D", 0.4);

        ProbabilityArrayInitializer initializer = new ProbabilityArrayInitializer();
        List<String> probabilityArray = initializer.initProbabilityArray(modelMap, 10000);

        // 打印初始化的概率数组
        HashMap<String, Integer> hashMap = new HashMap<>();
        for (String label : probabilityArray) {
            hashMap.put(label, hashMap.getOrDefault(label, 0) + 1);
        }
        System.out.println(hashMap);
    }
}
